package com.bw;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.dataproc.SplitBatchOp;
import com.alibaba.alink.operator.batch.dataproc.vector.VectorAssemblerBatchOp;
import com.alibaba.alink.operator.batch.evaluation.EvalRegressionBatchOp;
import com.alibaba.alink.operator.batch.feature.OneHotPredictBatchOp;
import com.alibaba.alink.operator.batch.feature.OneHotTrainBatchOp;
import com.alibaba.alink.operator.batch.regression.LassoRegPredictBatchOp;
import com.alibaba.alink.operator.batch.regression.LassoRegTrainBatchOp;
import com.alibaba.alink.operator.batch.sink.AkSinkBatchOp;
import com.alibaba.alink.operator.batch.source.CsvSourceBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.common.evaluation.RegressionMetrics;
import com.alibaba.alink.pipeline.regression.LassoRegression;
import com.alibaba.alink.pipeline.regression.LassoRegressionModel;
import com.alibaba.alink.pipeline.regression.RidgeRegression;
import com.alibaba.alink.pipeline.regression.RidgeRegressionModel;

public class Test1 {
    public static void main(String[] args) throws Exception {
        BatchOperator.setParallelism(1);
        // 声明列和目标值
        String schemaStr="f0 string,f1 double,f2 double,f3 double,f4 double,f5 double,f6 double,f7 double,label int";
        String label="label";

        // 读取数据源
        CsvSourceBatchOp csvSourceBatchOp = new CsvSourceBatchOp()
                .setFilePath("datafile/abalone.data")
                .setSchemaStr(schemaStr)
                .setFieldDelimiter(",");
        // 前10条
        csvSourceBatchOp.print(10);



        // 独热编码
        BatchOperator <?> one_hot = new OneHotTrainBatchOp().setSelectedCols("f0");
        BatchOperator <?> model = csvSourceBatchOp.link(one_hot);
//        model.lazyPrint(10);
        BatchOperator <?> predictor = new OneHotPredictBatchOp().setOutputCols("f0_new");
        BatchOperator<?> one_hot_result = predictor.linkFrom(model, csvSourceBatchOp);
//        one_hot_result.print();


        // 向量聚合
        String [] features_new=new String[]{"f0_new","f1","f2","f3","f4","f5","f6","f7"};
        BatchOperator<?> vec=new VectorAssemblerBatchOp()
                .setSelectedCols(features_new)
                .setOutputCol("vol")
                .linkFrom(one_hot_result);



        // 拆分数据集
        BatchOperator <?> spliter = new SplitBatchOp().setFraction(0.8);
        // 训练集
        BatchOperator<?> train_data = spliter.linkFrom(vec);
//        BatchOperator<?> train_data = spliter.lazyPrint(-1);
        BatchOperator<?> test_data = spliter.getSideOutput(0);


        // 最终数据集
        vec.print();
        System.out.println("train_data.count() = " + train_data.count());
        System.out.println("test_data.count() = " + test_data.count());



        // 训练模型
        LassoRegression lasso = new LassoRegression()
                .setVectorCol("vol")
                .setLambda(0.1)
                .setLabelCol(label)
                .setMaxIter(10)
                .setLambda(0.2)
                .setOptimMethod("SGD")
                .setPredictionCol("pred");
        LassoRegressionModel lassoRegressionModel = lasso.fit(train_data);

        BatchOperator<?> result = lassoRegressionModel.transform(test_data);
        result.print();

        // 评估
        RegressionMetrics metrics = new EvalRegressionBatchOp().setPredictionCol("pred").setLabelCol(label).linkFrom(
                result).collectMetrics();
        System.out.println("Total Samples Number:" + metrics.getCount());
        System.out.println("SSE:" + metrics.getSse());
        System.out.println("SAE:" + metrics.getSae());
        System.out.println("RMSE:" + metrics.getRmse());
        System.out.println("R2:" + metrics.getR2());
        System.out.println("MES:" + metrics.getMse());

        //-------------------------------------岭回归
        RidgeRegression ridge = new RidgeRegression()
//                .setFeatureCols()
                .setVectorCol("vol")
                .setLambda(0.1)
                .setLabelCol(label)
                .setMaxIter(10)
                .setOptimMethod("SGD")
                .setPredictionCol("pred");
        RidgeRegressionModel ridgeRegressionModel = ridge.fit(train_data);
        BatchOperator<?> result1 = ridgeRegressionModel.transform(test_data);


        result1.print();
        // 评估
        RegressionMetrics metrics1 = new EvalRegressionBatchOp().setPredictionCol("pred").setLabelCol(label).linkFrom(
                result1).collectMetrics();
        System.out.println("Total Samples Number:" + metrics1.getCount());
        System.out.println("SSE:" + metrics1.getSse());
        System.out.println("SAE:" + metrics1.getSae());
        System.out.println("RMSE:" + metrics1.getRmse());
        System.out.println("R2:" + metrics1.getR2());
        System.out.println("MES:" + metrics1.getMse());



        // 模型保存
//        AkSinkBatchOp csvSink1 = new AkSinkBatchOp();
//        csvSink1.setFilePath("E:\\Flink\\FlinkSql\\FlinkML\\model\\testGbdtRegPredictStreamOpModel");
//        result1.link(csvSink1);


    }
}
